
package edu.cmu.cs211.proximitycounter.counter;

import java.util.*;

/* This purely functional implementation of AVL trees is 
   adapted from set.ml in the standard ocaml distribution.
 
   Note that s.insert(x) returns a new set obtained by adding x to s,
   as does s.delete(x).  These operations do not change the set s,
   but instead return a new set modified accordingly.

   Danny Sleator, January 2011
*/

public class MyTreeSet<E extends Comparable<E>> implements Iterable<E> {
    /* Here's the public interface to this class: 
       constructors (all O(1) time):

         Set(E x)                   a new singleton set containing x
         Set()                      a new empty set
         Set(Set s)                 make a new copy of a set
       
       Operations (all O(log n) time unless otherwise specified):

         boolean isEmpty()          returns true if the set is empty.  O(1) time
         int size()                 returns the size of the set. O(n) time
         boolean contains(E x)      returns true if x is in the set
         Set<E> insert(E x)         adds x to the set
         Set<E> delete(E x)         delete x from the set if it is in the set
         E min()                    return the first (minimum) in the set
         E max()                    return the last (maximum) in the set
         Set tailSet(E x, boolean inclusive)   new set of all elements < x (<= if inclusive=true)
         Set headSet(E x, boolean inclusive)   new set of all elements > x (>= if inclusive=true)
         Set join(Set l, Set r)     a new set that is the union of these sets. all elements of l must be < all elements of r
         Iterator<E> iterator()     returns an iterator for the elements of this set
    */

    public static class Node<E> {
    	public final E data;
    	public final int height;
    	public final Node<E> left, right; 
    	
    	
    	Node(Node<E> l, E d, Node<E> r) {
    		left = l; data = d; right = r;
    		height = 1+Math.max (getHeight(left), getHeight(right));
    	}
    }

    public Node<E> root;

    /* This is the "balance factor" -- the maximum allowed height difference
     * between siblings.  Setting it to larger values, increases the worst-case
     * depth of the tree, but reduces the amount of rebalancing that happens.
    */
    static final int BALANCE_FACTOR=1;  

    // Constructors:
    public MyTreeSet(E x) 
    { 
    	root = new Node<E>(null, x, null); 
    }
    
    public MyTreeSet() 
    { 
    	root = null; 
    }
    
    /* copy a set */
    public MyTreeSet(MyTreeSet<E> s) 
    { 
    	root = s.root; 
    }
    
    /* Take a node and turn it into a set. */
    private MyTreeSet(Node<E> r) 
    { 
    	root = r; 
    } 

    private static int getHeight (Node<?> t) {
    	return (t==null)? 0: t.height;
    }


    /**
     * This method returns the number of keys in the range [a,b].
     */
    public int rangeCount(E a, E b) {
    	//throw new RuntimeException ("You need to implement this method");
    	if(root != null)
    		return visit(root, a, b, max(), min());
    	return 0;
    }
    
    /*
     * This method recursively counts all nodes that are within the range of a, b
     * This algorithm has a running time of O(logn)
     */
    public int visit(Node<E> n, E a, E b, E max, E min)
    {
    	// if node is null or if the range is unreasonable
    	if(n == null || a.compareTo(b)>0 || a.compareTo(max) > 0 || b.compareTo(min) < 0)
    		return 0;
    	// node we are looking at is too small, visit the right subtree
    	if(a.compareTo(n.data)>0 )
    	{
    		return visit(n.right, a, b, max, min);
    	}
    	// node we are looking at is too big, visit the left subtree
    	else if(b.compareTo(n.data)<0)
    	{
    		return visit(n.left, a, b, max, min);
    	}
    	// if the current node is a leaf: for speeding up purpose
    	else if(n.left == null && n.right == null)
    	{
    		return 1;
    	}
    	//node we are looking at is within the range, add 1 to count, and visit the left and right subtrees
    	else
    	{
    		return 1 + visit(n.left, a, b, max, min) + visit(n.right, a, b, max, min);
    	}
    }
    
    /**
     * This method returns true if the set is empty
     */
    public boolean isEmpty() {
    	return root == null;
    }

    /**
     * This method returns the number of elements in the set.
     * It runs in O(n) time for a set of size n.
     */
    public int size() {
    	return sizeAux(root);
    }
    
    private int sizeAux(Node<E> t) {
    	if(t==null) 
    		return 0;
    	
    	return sizeAux(t.left) + 1 + sizeAux(t.right);
    }
    
    /**
     * This method returns true if x in the set.
     */
    public boolean contains(E x) 
    {
    	return containsAux(x, root);
    }

    private boolean containsAux(E x, Node<E> t) {
    	if (t==null) 
    		return false;
	
    	int c = x.compareTo(t.data);
    	
    	if (c==0) 
    		return true;
    	if (c<0) 
    		return containsAux(x, t.left);
    	else 
    		return containsAux(x, t.right);
    }

    /**
     * This method returns the largest element in the set
     * that is < x.  If there is no such element, then it
     * returns null.
     */
    public E predecessor(E x) {
        return predAux(root, x, null);
    }

    private E predAux(Node<E> t, E x, E best) {
        if (t==null) 
        	return best;
        
        int c = x.compareTo(t.data);
        
        if (c>0) 
        	return predAux(t.right, x, t.data);
        else 
        	return predAux(t.left, x, best);
    }

    /**
     * This method returns the smallest element in the set
     * that is > x.  If there is no such element, then it
     * returns null.
     */
    public E successor(E x) {
        return succAux(root, x, null);
    }

    private E succAux(Node<E> t, E x, E best) {
        if (t==null) 
        	return best;
        
        int c = x.compareTo(t.data);
        
        if (c<0) 
        	return succAux(t.left, x, t.data);
        else 
        	return succAux(t.right, x, best);
    }

    
    /* Creates a new node with left child l, value v and right child r.
       We must have all elements of l < v < all elements of r.
       l and r must be balanced and | height(l) - height(r) | <= B. 
    */
    private Node<E> create(Node<E> l, E v, Node<E> r) {
    	return new Node<E>(l,v,r);
    }
    
    /* Same as create, but performs one step of rebalancing if necessary.
       Assumes l and r balanced and | height(l) - height(r) | <= B+1.
    */
    private Node<E> bal(Node<E> l, E v, Node<E> r) {
		int hl = getHeight(l);
		int hr = getHeight(r);
		
		if (hl > hr + BALANCE_FACTOR) {
		    if (l == null) 
		    	throw new NullPointerException("balance error");
		    
		    if (getHeight(l.left) >= getHeight(l.right)) {
		    	return create (l.left, l.data, create(l.right, v, r));
		    } 
		    else {
		    	if (l.right == null) 
		    		throw new NullPointerException();
		    	
		    	return create (create(l.left, l.data, l.right.left), l.right.data, create(l.right.right, v, r));
		    }
		}
		else if (hr > hl + BALANCE_FACTOR) {
		    if (getHeight(r.right) >= getHeight(r.left)) {
		    	return create(create(l,v,r.left), r.data, r.right);
		    } else {
			if (r.left==null) throw new RuntimeException("balance error");
				return create(create(l,v,r.left.left), r.left.data, create(r.left.right, r.data, r.right));
		    }
		} 
		else {
		    return create(l,v,r);
		}
    }
    
    /**
     * This method inserts x into the set.  Nothing happens
     * if x is already in the set.
     */
    public MyTreeSet<E> insert(E x) {
    	return new MyTreeSet<E> (insertAux (x, root));
    }

    private Node<E> insertAux (E x, Node<E> t) {
    	if (t == null) 
    		return create(null, x, null);
    	
    	int c = x.compareTo(t.data);
    	
    	if (c==0) 
    		return t;
    	if (c<0) 
    		return bal (insertAux(x, t.left), t.data, t.right);
    	else 
    		return bal (t.left, t.data, insertAux(x, t.right));
    }
    
    /**
     * This method returns the minimum element of the set.
     */
    public E min() {
    	if (root == null) 
    		throw new NoSuchElementException();
	
    	return minAux(root);
    }

    private E minAux(Node<E> t) {
    	if (t.left==null) 
    		return t.data;
    	return minAux(t.left);
    }

    /**
     * This method returns the maximum element of the set.
     */    
    public E max() {
    	if (root==null) 
    		throw new NoSuchElementException();
	
    	return maxAux(root);
    }

    private E maxAux(Node<E> t) {
    	if (t.right==null) 
    		return t.data;
    	
    	return maxAux(t.right);
    }

    private Node<E> deleteMin (Node<E> t) {
    	if (t==null) 
    		throw new NoSuchElementException();
    	
    	if (t.left==null) return t.right;
    	
    	return bal (deleteMin(t.left), t.data, t.right);
    }
    
    /* Merge two trees l and r into one.
       All elements of l must precede the elements of r.
       Assume | height l - height r | <= B.
    */
    
    private Node<E> merge(Node<E> t1, Node<E> t2) {
    	if (t1==null) return t2;
    	if (t2==null) return t1;
    	
    	return bal(t1, minAux(t2), deleteMin(t2));
    }

    /**
     * This method deletes x from the set.  Nothing happens if x
     * is not in the set.
     */

    public MyTreeSet<E> delete(E x) {
    	return new MyTreeSet<E>(deleteAux (x, root));
    }

    private Node<E> deleteAux (E x, Node<E> t) {
    	if (t==null) 
    		return null;
    	
    	int c = x.compareTo(t.data);
    	
    	if (c==0) 
    		return merge(t.left, t.right);
    	if (c<0) 
    		return bal (deleteAux(x,t.left), t.data, t.right);
    	else 
    		return bal(t.left, t.data, deleteAux(x, t.right));
    }
    
    /* the following two methods are for joining two sets, where
       the first is less than the second */

    /**
     * This method joins two sets together.  All the elements
     * of the first must be less than all the elements of the second.
     * Throws an exception if this is not the case.
     */
    public MyTreeSet<E> join(MyTreeSet<E> s2) {
    	return new MyTreeSet<E>(joinAux(root, s2.root));
    }
    
    private Node<E> joinAux(Node<E> t1, Node<E> t2) {
    	if (t1==null) return t2;
    	if (t2==null) return t1;
	
    	if (maxAux(t1).compareTo(minAux(t2)) >=0) 
    		throw new RuntimeException("improper join");
    	
    	return join3 (t1, minAux(t2), deleteMin(t2));
    }

    /* Same as create and bal, but no assumptions are made on the
       relative heights of l and r. */
    
    private Node<E> join3 (Node<E> l, E v, Node<E> r) {
    	if (l==null) return insertAux(v,r);
    	if (r==null) return insertAux(v,l);
    	
    	if (l.height > r.height + BALANCE_FACTOR) return bal(l.left, l.data, join3(l.right, v, r));
    	if (r.height > l.height + BALANCE_FACTOR) return bal(join3(l, v, r.left), r.data, r.right);
    	return create(l,v,r);
    }

    private class SplitReturn {
    	Node<E> l, r;
    	boolean found;
    }
    
    private SplitReturn splitAux (E x, Node<E> t, SplitReturn sr) {
		if (t==null) {
		    sr.l = sr.r = null;
		    sr.found = false;
		} else {
		    int c = x.compareTo(t.data);
		    if (c==0) {
				sr.l = t.left;
				sr.r = t.right;
				sr.found = true;
		    } else if (c<0) {
				sr = splitAux(x, t.left, sr);
				sr.r = join3(sr.r, t.data, t.right);
		    } else {
				sr = splitAux(x, t.right, sr);
				sr.l = join3(t.left, t.data, sr.l);
		    }
		}

		return sr;
    }
    
    private MyTreeSet<E> half_split(E x, boolean inclusive, boolean takehead) {
		SplitReturn sr = new SplitReturn();  /* where we put the return values from split */
		
		sr = splitAux(x, root, sr);
		
		MyTreeSet<E> hs = new MyTreeSet<E>((takehead)?sr.l:sr.r);
		
		if (inclusive && sr.found) return hs.insert(x);
		return hs;
    }

    /**
     * This returns a subset of the given set containing all elements
     * that are greater than x (greater than or equal to if inclusive
     * is true).
     */
    public MyTreeSet<E> tailSet(E x, boolean inclusive) {
    	return half_split(x, inclusive, !false);
    }
    
    /**
     * This returns a subset of the given set containing all elements
     * that are less than x (less than or equal to if inclusive
     * is true).
     */
    public MyTreeSet<E> headSet(E x, boolean inclusive) {
    	return half_split(x, inclusive, !true);
    }
    
    public Iterator<E> iterator () { 
		final MyTreeSet<E> s = new MyTreeSet<E>(root);
	
		return new Iterator<E>() {
		    public boolean hasNext() { return !s.isEmpty(); }
		    
		    public E next() { 
			E c = s.min();
			s.root = deleteAux(c, s.root); 
			return c; 
		    } 
		    
		    public void remove() { throw new UnsupportedOperationException(); }
		};
    }
    
    private String myToString(Node<E> t) {
		if (t == null) return "";
		return "("+myToString(t.left) + " " + t.data + " " + myToString(t.right)+")"; 
	}
    
    public String toString() {
    	return myToString(root);
    }

    private String myurl(Node<E> t) {
    	if (t == null) return "";
    	return t.data+"."+myurl(t.left)+myurl(t.right);
    }
    
    private String url() {
    	return "http://www.link.cs.cmu.edu/cgi-bin/splay/splay-cgi.pl/"+myurl(root)+"/";
    }

    /* return the height of the tree, or -1 if not balanced */
    int checkHeight(Node<E> t) {  
    	if (t==null) return 0;
    	int lb = checkHeight(t.left);
    	int rb = checkHeight(t.right);
    	if (lb<0 || rb<0) return -1;
    	if (Math.abs(lb-rb) > BALANCE_FACTOR) return -1;
    	return 1+Math.max(lb,rb);
    }
    
    boolean checkBalance() {  
    	return checkHeight(root) < 0;
    }
}
